from bEandI_classes import *
from general_neuro.conductance_neuron import Conductance_neuron
from general_neuro.alpha_synapse import Alpha_synapse
from bEandI_classes import Cell as NEURON_Cell
from scipy.integrate import odeint
from numpy import array

class Cell(NEURON_Cell):
    """ This is just like the Cell class in bEandI_classes, except that it uses python (scipy) for
        all the simulations. """
        
    def exercise(self):
        '''create the model and run the simulation'''
        # the results of a simulation
        self.sim_results = {}

        # set up a list to hold synapse current vectors.
        self.svects_list = range(len(self.synapses))

        self.model = range(len(self.data))
        self.sv_results = range(len(self.data))
        for i in range(len(self.data)): # loop through each holding potential
            hp = self.data[i].holding_potential
            # for each holding potential, we'll have simulation results
            self.sim_results[hp] = {}
            # --------------------- create the model. ---------------------------
            self.model[i] = model = Conductance_neuron(g_leak=self.g_pas, E_leak = hp,\
                                     tau = self.tau, V_init = hp)

            # set up the synapses
            for syn in self.synapses:
                onset = syn.onset
                tau   = syn.tau
                gmax  = syn.gmax
                E_rev = syn.rev
                name  = syn.name
                model.register_input(name,Alpha_synapse(onset=onset, tau=tau,\
                                  gmax=gmax, E_rev=E_rev))
            # run the model
            sv_init = model.get_sv_init_list()
            self.sv_results[i] = sv_results = odeint(model.dX_dt, sv_init, self.t)
            # the voltage is always the first element in the sv_results.T
            self.sim_results[hp]['V_trace'] = sv_results.T[0]

    def calculate_currents(self):
        for i in range(len(self.data)): # loop through each holding potential
            hp = self.data[i].holding_potential
            # recreate the current traces
            model = self.model[i]
            sv_results = self.sv_results[i]

            currents = []
            for sv,t in zip(sv_results,self.t):
                currents.append(model.calculate_currents(sv,t))
            currents = array(currents)

            index = 0
            for name in model.inputs.iterkeys():
                key = "I_" + name + "_trace"
                self.sim_results[hp][key] = currents.T[index]
                index += 1
            for name in model.currents.iterkeys():
                key = "I_" + name + "_trace"
                self.sim_results[hp][key] = currents.T[index]
                index += 1
                
                      
